import cv2
import torch
import numpy as np
import os
import random


class DataLoader1:
    trLength = 0
    teLength = 0
    trInput = []
    teInput = []
    trOutput = []
    teOutput = []
    trPos = 0
    tePos = 0

    # 初始化加载
    def __init__(self, root, width=640, height=480, channal=3, trainRate=0.8):
        # 初始化路径
        if root[-1] != "/":
            root += "/"
        inpPath = root + "inputs/"
        labelsPath = root + "labels/"

        trainTxtPath = root + "train.txt"
        testTxtPath = root + "test.txt"

        # 读取数据
        print(">>>开始读取数据")
        ############################ 读取训练集
        cnt = 0
        lis = []
        with open(trainTxtPath, "r") as f:
            for line in f:
                lis.append(line[:-1])
        for f in lis:
            cnt += 1
            # if cnt>20:break
            # print(cnt,'in',len(lis),':',f)
            # 确定文件类型
            name, tp = f.split(".")
            if tp != "jpg" and tp != "png" and tp != "jpeg":
                continue
            inp = cv2.imread(inpPath + f)
            h, w, _ = inp.shape
            label = self.getLabels(labelsPath + name + ".txt")
            outp = self.createOutputImg(h, w, label)

            inp = self.fixInput(inp, width, height)
            outp = self.fixOutput(outp, width, height)

            self.trInput.append(inp)
            self.trOutput.append(outp)
        ############################ 读取测试集
        cnt = 0
        lis = []
        with open(testTxtPath, "r") as f:
            for line in f:
                lis.append(line[:-1])
        for f in lis:
            cnt += 1
            # if cnt>20:break
            # 确定文件类型
            name, tp = f.split(".")
            if tp != "jpg" and tp != "png" and tp != "jpeg":
                continue
            inp = cv2.imread(inpPath + f)
            h, w, _ = inp.shape
            label = self.getLabels(labelsPath + name + ".txt")
            outp = self.createOutputImg(h, w, label)

            inp = self.fixInput(inp, width, height)
            outp = self.fixOutput(outp, width, height)

            self.teInput.append(inp)
            self.teOutput.append(outp)

        self.trLength = len(self.trInput)
        self.teLength = len(self.teInput)
        print(">>>读取完毕")
        # 打乱顺序
        self.sort(self.trInput, self.trOutput)
        self.sort(self.teInput, self.teOutput)
        # 合并
        self.trInput = torch.cat(self.trInput, dim=0)
        self.trOutput = torch.cat(self.trOutput, dim=0)
        self.teInput = torch.cat(self.teInput, dim=0)
        self.teOutput = torch.cat(self.teOutput, dim=0)

        print("数据准备就绪")

    # 把输入图片转化成torch张量
    def fixInput(self, img, w=640, h=480):
        # 调节大小
        c = 3
        p = img.copy()
        r = np.zeros((h, w, c))
        hh, ww, cc = img.shape
        if ww / hh > w / h:
            k = int(hh * w / ww)
            margin = (h - k) // 2
            p = cv2.resize(p, (w, k))
            r[margin : margin + k] = p
        else:
            k = int(ww * h / hh)
            margin = (w - k) // 2
            p = cv2.resize(p, (k, h))
            r[:, margin : margin + k] = p
        r = torch.FloatTensor(r.T) / 255 * 2 - 1
        return torch.reshape(r, (1, c, w, h))

    # 把输出转化成张量
    def fixOutput(self, img, w=640, h=480):
        # 调节大小
        p = img.copy()
        r = np.zeros((h, w))
        hh, ww = img.shape
        if ww / hh > w / h:
            k = int(hh * w / ww)
            margin = (h - k) // 2
            p = cv2.resize(p, (w, k))
            r[margin : margin + k] = p
        else:
            k = int(ww * h / hh)
            margin = (w - k) // 2
            p = cv2.resize(p, (k, h))
            r[:, margin : margin + k] = p
        r = torch.FloatTensor(r.T) / 255
        return torch.reshape(r, (1, w, h))

    def getLabels(self, path):
        labels = []
        with open(path) as f:
            for line in f:
                labels.append(line.strip())
        lis = []
        for i in range(len(labels)):
            if len(labels[i]) < 5:
                continue
            l = labels[i].split()
            l = [int(l[0]), float(l[1]), float(l[2]), float(l[3]), float(l[4])]
            lis.append(l)
        return lis

    # 生成输出
    def createOutputImg(self, h, w, labels):
        img = np.zeros((h, w), dtype=np.uint8)
        for l in labels:
            if l[0] != 1 and l[0] != 2:
                continue
            x1, y1 = l[1] - l[3] / 2, l[2] - l[4] / 2
            x2, y2 = l[1] + l[3] / 2, l[2] + l[4] / 2
            x1, y1, x2, y2 = int(w * x1), int(h * y1), int(w * x2), int(h * y2)
            cv2.rectangle(img, (x1, y1), (x2, y2), 255, -1)
        return img

    # 打乱顺序
    def sort(self, lisIn, lisOut, m=0.7):
        n = len(lisIn)
        for k in range(int(n * m)):
            i = random.randint(0, n - 1)
            j = random.randint(0, n - 1)
            lisIn[i], lisIn[j] = lisIn[j], lisIn[i]
            lisOut[i], lisOut[j] = lisOut[j], lisOut[i]

    # 获取训练数据
    def getTrain(self, length):
        end = min(self.trPos + length, self.trLength)
        if end == self.trLength:
            self.trPos = self.trLength - length
        r = self.trInput[self.trPos : end], self.trOutput[self.trPos : end]
        self.trPos += length
        if self.trPos >= self.trLength:
            self.trPos = 0
        return r

    # 获取测试数据
    def getTest(self, length):
        end = min(self.tePos + length, self.teLength)
        if end == self.teLength:
            self.tePos = self.teLength - length
        r = self.teInput[self.tePos : end], self.teOutput[self.tePos : end]
        self.tePos += length
        if self.tePos >= self.teLength:
            self.tePos = 0
        return r


class DataLoader:
    trLength = 0
    teLength = 0
    trInput = []
    teInput = []
    trOutput = []
    teOutput = []
    trPos = 0
    tePos = 0

    # 初始化加载
    def __init__(self, root, width=640, height=480, channal=3, trainRate=0.8):
        # 初始化路径
        if root[-1] != "/":
            root += "/"
        inpPath = root + "inputs/"
        labelsPath = root + "labels/"

        trainTxtPath = root + "train.txt"
        testTxtPath = root + "test.txt"

        # 读取数据
        print(">>>开始读取数据")
        ############################ 读取训练集
        cnt = 0
        lis = []
        with open(trainTxtPath, "r") as f:
            for line in f:
                lis.append(line[:-1])
        for f in lis:
            cnt += 1
            # if cnt>20:break
            # print(cnt,'in',len(lis),':',f)
            # 确定文件类型
            name, tp = f.split(".")
            if tp != "jpg" and tp != "png" and tp != "jpeg":
                continue
            inp = cv2.imread(inpPath + f)
            outp = self.getLabels(labelsPath + name + ".txt")
            inp = self.fixInput(inp, width, height)

            self.trInput.append(inp)
            self.trOutput.append(outp)
        ############################ 读取测试集
        cnt = 0
        lis = []
        with open(testTxtPath, "r") as f:
            for line in f:
                lis.append(line[:-1])
        for f in lis:
            cnt += 1
            # if cnt>20:break
            # 确定文件类型
            name, tp = f.split(".")
            if tp != "jpg" and tp != "png" and tp != "jpeg":
                continue
            inp = cv2.imread(inpPath + f)
            outp = self.getLabels(labelsPath + name + ".txt")
            inp = self.fixInput(inp, width, height)

            self.teInput.append(inp)
            self.teOutput.append(outp)

        self.trLength = len(self.trInput)
        self.teLength = len(self.teInput)
        print(">>>读取完毕")
        # 打乱顺序
        self.sort(self.trInput, self.trOutput)
        self.sort(self.teInput, self.teOutput)
        # 合并
        self.trInput = torch.cat(self.trInput, dim=0)
        self.trOutput = torch.cat(self.trOutput, dim=0)
        self.teInput = torch.cat(self.teInput, dim=0)
        self.teOutput = torch.cat(self.teOutput, dim=0)

        print("数据准备就绪")

    # 把输入图片转化成torch张量
    def fixInput(self, img, w=640, h=480):
        # 调节大小
        c = 3
        p = img.copy()
        r = np.zeros((h, w, c))
        hh, ww, cc = img.shape
        if ww / hh > w / h:
            k = int(hh * w / ww)
            margin = (h - k) // 2
            p = cv2.resize(p, (w, k))
            r[margin : margin + k] = p
        else:
            k = int(ww * h / hh)
            margin = (w - k) // 2
            p = cv2.resize(p, (k, h))
            r[:, margin : margin + k] = p
        r = torch.FloatTensor(r.T) / 255 * 2 - 1
        return torch.reshape(r, (1, c, w, h))

    def getLabels(self, path):
        labels = []
        with open(path) as f:
            for line in f:
                labels.append(line.strip())
        lis = []
        for i in range(len(labels)):
            if len(labels[i]) < 5:
                continue
            l = labels[i].split()
            l = [int(l[0]), float(l[1]), float(l[2]), float(l[3]), float(l[4])]
            if l[0] == 1 or l[0] == 2:
                l[0] = 1
                lis.append(l)
        lis = torch.FloatTensor(l).view(1, -1, 5)
        return lis

    # 打乱顺序
    def sort(self, lisIn, lisOut, m=0.7):
        n = len(lisIn)
        for k in range(int(n * m)):
            i = random.randint(0, n - 1)
            j = random.randint(0, n - 1)
            lisIn[i], lisIn[j] = lisIn[j], lisIn[i]
            lisOut[i], lisOut[j] = lisOut[j], lisOut[i]

    # 获取训练数据
    def getTrain(self, length):
        end = min(self.trPos + length, self.trLength)
        if end == self.trLength:
            self.trPos = self.trLength - length
        r = self.trInput[self.trPos : end], self.trOutput[self.trPos : end]
        self.trPos += length
        if self.trPos >= self.trLength:
            self.trPos = 0
        return r

    # 获取测试数据
    def getTest(self, length):
        end = min(self.tePos + length, self.teLength)
        if end == self.teLength:
            self.tePos = self.teLength - length
        r = self.teInput[self.tePos : end], self.teOutput[self.tePos : end]
        self.tePos += length
        if self.tePos >= self.teLength:
            self.tePos = 0
        return r


if __name__ == "__main__":
    path = r"./dataset"
    dataLoader = DataLoader(path, width=128, height=128)
    inp, outp = dataLoader.getTest(4)
    print(inp.shape, outp.shape)
